{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 载入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# import libs\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection import FasterRCNN\n",
    "from torchvision.models.detection.rpn import AnchorGenerator\n",
    "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
    "from torchvision.models.detection import MaskRCNN\n",
    "import torchvision.transforms as tT\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image, ExifTags\n",
    "import json\n",
    "import gradio as gr\n",
    "import PySimpleGUI as sg\n",
    "\n",
    "sys.path.insert(0, './torchvision_det_references') #确保可以通过下面的语句导入位于子目录中的包\n",
    "import utils\n",
    "from engine import train_one_epoch, evaluate\n",
    "import transforms as T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 检查标注的正确性，把错误标注写入文件\n",
    "TODO:目前还检查不出如下错误：指针和零刻度方向的点标记顺序错误。比如假设第一个点是指针的旋转中心，第二个点是指针尖端；假设零刻度方向的第一个点是0，第二个点是5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# folder = \"D:\\\\Downloads\\\\DATA\\\\水表\\\\标注\\\\8p\\\\8P blur\"\n",
    "# folder = \"D:\\\\Downloads\\\\DATA\\\\水表\\\\标注\\\\3p\\\\3P uneven\"\n",
    "# folder = \"D:\\\\DATA\\\\multi-pointer-meter\\\\pointer-8\\\\8p even\"\n",
    "# folder = \"D:\\\\DATA\\\\multi-pointer-meter\\\\pointer-3\\\\3P even\"\n",
    "# folder = \"D:\\\\Downloads\\\\DATA\\\\meters\\\\dataset-10-29-jpg\\\\dataset\\\\4PEven\"\n",
    "folder = \"C:\\\\Users\\\\86198\\\\Desktop\\\\8针水表\\\\8针水表\\\\8P Uneven\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# check the annotaion files and write checking results to the local file ann_err.txt\n",
    "\n",
    "files = list(sorted(os.listdir(folder)))\n",
    "im_files = [f for f in files if f.endswith(\"jpg\")]\n",
    "print(\"总共有：\"+str(len(im_files))+\"张jpg格式的图像，应该是水表图像。\")\n",
    "\n",
    "no_jpg_json_files = []\n",
    "no_png_json_files = []\n",
    "no_png_mask_files = []\n",
    "\n",
    "with open(os.path.join(folder, \"err_files.txt\"), 'w') as err_file:\n",
    "    for f in im_files:\n",
    "        err_lst = [] #保存针对该图像文件的错误信息\n",
    "        # 检查是否缺基本标注文件和掩膜标注或缺掩膜图像文件\n",
    "        if f+\".json\" not in files:\n",
    "            no_jpg_json_files.append(f)\n",
    "            print(f+\"缺少对应的标注文件\")\n",
    "            err_lst.append(\"缺少对应的标注文件\")\n",
    "            if len(no_jpg_json_files)==1:\n",
    "                print(\"存在缺少基本标注json的文件\")           \n",
    "\n",
    "        else:\n",
    "\n",
    "            # if f[:-4]+\".png\" not in files:\n",
    "            #     no_png_mask_files.append(f)\n",
    "            #     err_lst.append(\"缺少对应的轮廓掩膜Png图像\")\n",
    "            #     print(f + \"缺少对应的轮廓掩膜Png图像\")\n",
    "            #     if len(no_png_mask_files)==1:\n",
    "            #         print(\"\\n存在缺少轮廓掩膜png的文件\")\n",
    "\n",
    "            # if f[:-4]+\".png\"+\".json\" not in files:\n",
    "            #     err_lst.append(\"缺少轮廓掩膜json的文件\")\n",
    "            #     no_png_json_files.append(f)\n",
    "            #     if len(no_png_json_files)==1:\n",
    "            #         print(\"\\n存在缺少轮廓掩膜json的文件\")\n",
    "\n",
    "\n",
    "            #读取基本标记和掩膜信息，并检查细节的合法性    \n",
    "            file_path = os.path.join(folder, f+\".json\")\n",
    "            with open(file_path, 'r', encoding='utf-8') as fjson:\n",
    "                injson = json.load(fjson)\n",
    "\n",
    "                #3针的标注方式与8针的不同，多了一个step，估计是表盘轮廓\n",
    "            if_4_steps = \"step_4\" in injson.keys()\n",
    "            if if_4_steps:\n",
    "                key_bboxes = \"step_2\"\n",
    "                key_zero_scale_dir = \"step_3\"\n",
    "                key_pointers_dir = \"step_4\"\n",
    "            else:\n",
    "                key_bboxes = \"step_1\"\n",
    "                key_zero_scale_dir = \"step_2\"\n",
    "                key_pointers_dir = \"step_3\"\n",
    "                \n",
    "                #读取指针的包围盒\n",
    "            n_pointers = len(injson[key_bboxes]['result'])\n",
    "            if (n_pointers!=3 and n_pointers!=8 and n_pointers!=4):\n",
    "                err_lst.append(\"指针数目量不是3也不是8\")    \n",
    "\n",
    "                #读取指针方向\n",
    "            if f=='IMG_2852_00.jpg':\n",
    "                tmp=1\n",
    "            n_pointers2 = len(injson[key_pointers_dir]['result'])\n",
    "            if(n_pointers!=n_pointers2):\n",
    "                err_lst.append(\"包围盒的数量与指针方向标注的数量不一致\")\n",
    "\n",
    "                #检查指针方向的两点是否距离过近                      \n",
    "            for line in injson[key_pointers_dir]['result']:\n",
    "                pt_lst = line['pointList']      \n",
    "                line_vec = np.array([pt_lst[0]['x']-pt_lst[1]['x'], pt_lst[0]['y']-pt_lst[1]['y']] )                \n",
    "                if np.linalg.norm(line_vec)<20:\n",
    "                    err_lst.append(\"指针方向的两点距离过近\")\n",
    "                    break\n",
    "\n",
    "                #检查指针是否在包围盒内     \n",
    "                # ！！注意：一定要先讲包围盒列表和指针标注列表按id进行排序，避免没有对齐的问题。 \n",
    "            injson[key_bboxes]['result'].sort(key=lambda x:x[\"id\"])   \n",
    "            injson[key_pointers_dir]['result'].sort(key=lambda x:x[\"sourceID\"])   \n",
    "            if(n_pointers==n_pointers2):\n",
    "                for i in range(n_pointers):\n",
    "                    bbox = injson[key_bboxes]['result'][i] # dict\n",
    "                    pointer = injson[key_pointers_dir]['result'][i]['pointList'] # dict\n",
    "                    if_in_bbox = False\n",
    "                    for j in range(2):\n",
    "                        pt = pointer[j]\n",
    "                        if 0<pt['x']-bbox['x']<=bbox['width'] and 0<pt['y']-bbox['y']<=bbox['height']:\n",
    "                            if_in_bbox = True\n",
    "                            break\n",
    "                    if if_in_bbox == False:\n",
    "                        err_lst.append(\"指针方向的标注没有位于对应的包围盒内\")\n",
    "                        break\n",
    "\n",
    "\n",
    "                #检查零刻度方向的两点是否距离过近\n",
    "            n_zero_scale_dir = len(injson[key_zero_scale_dir]['result'])            \n",
    "            for line in injson[key_zero_scale_dir]['result']:\n",
    "                pt_lst = line['pointList']      \n",
    "                line_vec = np.array([pt_lst[0]['x']-pt_lst[1]['x'], pt_lst[0]['y']-pt_lst[1]['y']] )\n",
    "                if np.linalg.norm(line_vec)<20:\n",
    "                    err_lst.append(\"零刻度方向的两点距离过近\")\n",
    "                    break\n",
    "\n",
    "                #zero_scale_dir_lst.append([pt_lst[0]['x'],pt_lst[0]['y'],pt_lst[1]['x'],pt_lst[1]['y']]) # [x1,y1,x2,y2]第1个点是起点,据说是从刻度0到5\n",
    "\n",
    "\n",
    "        #如有错误，则把该图像所包含的错误信息写入文件\n",
    "        if len(err_lst)>0:\n",
    "            err_file.write('\\n')\n",
    "            err_file.write(f+\":\")\n",
    "            for err in err_lst:\n",
    "                err_file.write(\"||\"+err)\n",
    "\n",
    "print(\"\\ndone\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 浏览标注结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "#define loading functions\n",
    "def load_valid_imfile_names(folder:str)->list:\n",
    "    files = list(sorted(os.listdir(folder)))\n",
    "    im_files = [f for f in files if f.endswith(\"jpg\")]\n",
    "    #filter out the illegal files \n",
    "    assert(os.path.exists(os.path.join(folder, \"err_files.txt\")))\n",
    "    illegel_lst = []\n",
    "    with open(os.path.join(folder, \"err_files.txt\"),'r') as err_file:\n",
    "        lines = err_file.readlines()\n",
    "        for line in lines:\n",
    "            if len(line)>0:\n",
    "                pos_comma=line.find(':')\n",
    "                if pos_comma>-1:\n",
    "                    illegel_lst.append(line[:pos_comma])\n",
    "    \n",
    "    legal_lst = []\n",
    "    for f in im_files:\n",
    "        if f not in illegel_lst:\n",
    "            legal_lst.append(f)\n",
    "    \n",
    "    return legal_lst\n",
    "\n",
    "def load_anno(folder:str, im_fname:str)->dict:\n",
    "    #读取标注的json文件内容\n",
    "    file_path = os.path.join(folder, im_fname+\".json\")\n",
    "    with open(file_path, 'r', encoding='utf-8') as f:\n",
    "        injson = json.load(f)\n",
    "    #3针的标注方式与8针的不同，多了一个step，估计是表盘轮廓\n",
    "    if_4_steps = \"step_4\" in injson.keys()\n",
    "    if if_4_steps:\n",
    "        key_bboxes = \"step_2\"\n",
    "        key_zero_scale_dir = \"step_3\"\n",
    "        key_pointers_dir = \"step_4\"\n",
    "    else:\n",
    "        key_bboxes = \"step_1\"\n",
    "        key_zero_scale_dir = \"step_2\"\n",
    "        key_pointers_dir = \"step_3\"\n",
    "    \n",
    "    # ！！注意：一定要先讲包围盒列表和指针标注列表按id进行排序，避免没有对齐的问题。 \n",
    "    injson[key_bboxes]['result'].sort(key=lambda x:x[\"id\"])   \n",
    "    injson[key_pointers_dir]['result'].sort(key=lambda x:x[\"sourceID\"])   #它的sourceID对应包围盒的id\n",
    "\n",
    "    #读取包围盒信息\n",
    "    n_pointers = len(injson[key_bboxes]['result'])\n",
    "    bboxes_lst = []    \n",
    "    for bbox in injson[key_bboxes]['result']:\n",
    "        bboxes_lst.append([bbox['x'],bbox['y'],bbox['width'],bbox['height']])        \n",
    "    \n",
    "    #读取零刻度方向  \n",
    "    n_zero_scale_dir = len(injson[key_zero_scale_dir]['result'])\n",
    "    zero_scale_dir_lst = []\n",
    "    for line in injson[key_zero_scale_dir]['result']:\n",
    "        pt_lst = line['pointList']        \n",
    "        zero_scale_dir_lst.append([pt_lst[0]['x'],pt_lst[0]['y'],pt_lst[1]['x'],pt_lst[1]['y']]) # [x1,y1,x2,y2]第1个点是起点,据说是从刻度0到5\n",
    "        \n",
    "    #读取指针方向\n",
    "    n_pointers2 = len(injson[key_pointers_dir]['result'])   \n",
    "    pointer_dir_lst = [] \n",
    "    for line in injson[key_pointers_dir]['result']:\n",
    "        pt_lst = line['pointList']\n",
    "        if n_pointers2==8:\n",
    "            pointer_dir_lst.append([pt_lst[0]['x'],pt_lst[0]['y'],pt_lst[1]['x'],pt_lst[1]['y']]) # [x1,y1,x2,y2]第1个点是起点，是旋转中心\n",
    "        elif n_pointers2==3:\n",
    "            pointer_dir_lst.append([pt_lst[1]['x'],pt_lst[1]['y'],pt_lst[0]['x'],pt_lst[0]['y']]) # [x1,y1,x2,y2]第1个点是起点，是旋转中心\n",
    "    #读取表盘轮廓  TODO\n",
    "\n",
    "    anno={\"bbox\":bboxes_lst, \"zdir\":zero_scale_dir_lst, \"pdir\":pointer_dir_lst}\n",
    "    return anno\n",
    "\n",
    "def rotate_im_accord_exiftag(im:Image)->Image:\n",
    "    #处理手机拍照时的旋转问题(ref: https://blog.csdn.net/qq_34904125/article/details/120918232)\n",
    "    for orientation in ExifTags.TAGS.keys():\n",
    "        if ExifTags.TAGS[orientation]=='Orientation':\n",
    "            break\n",
    "    exif = dict(im._getexif().items())\n",
    "    if orientation in exif:       \n",
    "        if exif[orientation] == 3:\n",
    "            im = im.rotate(180,expand=True)\n",
    "        elif exif[orientation]==6:\n",
    "            im = im.rotate(270,expand=True)\n",
    "        elif exif[orientation]==8:\n",
    "            im=im.rotate(90, expand=True)\n",
    "    \n",
    "    return im\n",
    "\n",
    "def draw_line(x1,y1,x2,y2,line_color, start_color, end_color):\n",
    "    plt.plot([x1,x2],[y1,y2], color=line_color)\n",
    "    plt.plot(x1,y1, '*', color=start_color)\n",
    "    plt.plot(x2,y2,'*', color=end_color)\n",
    "\n",
    "\n",
    "def draw_im_with_anno(folder:str, im_fname:str, ann:dict):\n",
    "    fpath = os.path.join(folder, im_fname)\n",
    "    im = Image.open(fpath)\n",
    "    im = rotate_im_accord_exiftag(im)\n",
    "        \n",
    "    #显示图像\n",
    "    plt.imshow(im)\n",
    "\n",
    "    color_lst = [\"red\",\"green\", \"gold\", \"orange\",  \"lime\", \"darkred\",\"deepskyblue\", \"violet\"]\n",
    "\n",
    "    #绘制包围盒\n",
    "    n_bbox = len(ann[\"bbox\"])\n",
    "    for i in range(n_bbox):\n",
    "        bbox=ann['bbox'][i]\n",
    "        plt.gca().add_patch(plt.Rectangle((bbox[0],bbox[1]), bbox[2],bbox[3],fill=False, color=color_lst[i],linewidth=1))\n",
    "\n",
    "    #绘制指针方向\n",
    "    for i in range(n_bbox):\n",
    "        pts = ann['pdir'][i]\n",
    "        draw_line(pts[0],pts[1],pts[2],pts[3],color_lst[i], \"red\", 'blue')\n",
    "\n",
    "    #绘制零刻度方向\n",
    "    n_zeros = len(ann['zdir'])\n",
    "    for i in range(n_zeros):\n",
    "        pts=ann['zdir'][i]\n",
    "        draw_line(pts[0],pts[1],pts[2],pts[3],color_lst[i], \"red\", 'blue')\n",
    "\n",
    "\n",
    "flst = load_valid_imfile_names(folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# # show with gradio\n",
    "# def show_img_anno(x):\n",
    "#     fid = int(x)\n",
    "#     anno = load_anno(folder, flst[fid])\n",
    "#     fig = plt.figure()\n",
    "#     draw_im_with_anno(folder, flst[fid], anno)\n",
    "#     return fig    \n",
    "\n",
    "# def flip_image(x):\n",
    "#     return np.fliplr(x)\n",
    "\n",
    "# with gr.Blocks() as demo:\n",
    "#     gr.Markdown(\"显示水表的标注信息.\")\n",
    "#     with gr.Tabs():\n",
    "#         with gr.TabItem(\"指针和零刻度的信息\"):\n",
    "#             file_id = gr.Slider(0, len(flst)-1,step=1)\n",
    "#             im_show = gr.Plot()\n",
    "#             show_button = gr.Button(\"显示\")\n",
    "#         with gr.TabItem(\"表盘区域的分割图\"):\n",
    "#             with gr.Row():\n",
    "#                 image_input = gr.Image()\n",
    "#                 image_output = gr.Image()\n",
    "#             image_button = gr.Button(\"Flip\")    \n",
    "    \n",
    "#     show_button.click(show_img_anno, inputs=file_id, outputs=im_show)\n",
    "#     image_button.click(flip_image, inputs=image_input, outputs=image_output)\n",
    "    \n",
    "# demo.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# show with pysimplegui\n",
    "%matplotlib auto\n",
    "\n",
    "layout = [  [sg.Text('共有'+str(len(flst))+\"个带标注的水表图像\")],\n",
    "            [sg.Text('请输入要查看的图像序号(从0开始):'), sg.InputText('0')],\n",
    "            [sg.Checkbox('是否新开窗口查看', k='b_new_window')],\n",
    "            [sg.Button('Ok'), sg.Button('Cancel')] ]\n",
    "\n",
    "# Create the Window\n",
    "window = sg.Window('水表标注查看程序', layout)\n",
    "fig = plt.figure()\n",
    "# Event Loop to process \"events\" and get the \"values\" of the inputs\n",
    "while True:\n",
    "    event, values = window.read()\n",
    "    if event == sg.WIN_CLOSED or event == 'Cancel': # if user closes window or clicks cancel\n",
    "        break\n",
    "    if(int(values[0])>=0):\n",
    "        fid = int(values[0])\n",
    "        anno = load_anno(folder, flst[fid])\n",
    "        if values['b_new_window']:\n",
    "            fig = plt.figure()\n",
    "        plt.cla()\n",
    "        draw_im_with_anno(folder, flst[fid], anno)\n",
    "        plt.draw()\n",
    "\n",
    "window.close()\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义数据载入类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# create a dataloader class\n",
    "class MultiPointerMeterDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, root, transforms, down_scale_factor=6):\n",
    "        self.root = root # root folder path\n",
    "        self.transforms = transforms # data transformations\n",
    "        self.down_scale_factor = down_scale_factor # 图像缩小为原图的比例\n",
    "                            \n",
    "        self.imgs = load_valid_imfile_names(root)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        fpath = os.path.join(self.root, self.imgs[idx])\n",
    "        img = Image.open(fpath)\n",
    "        img = rotate_im_accord_exiftag(img)\n",
    "        img = img.convert('RGB')\n",
    "        #缩小图像\n",
    "        im_sz=(img.size[0]//self.down_scale_factor, img.size[1]//self.down_scale_factor)\n",
    "        img = img.resize(im_sz)\n",
    "\n",
    "        #载入标注信息\n",
    "        anno = load_anno(self.root, self.imgs[idx])\n",
    "        #获得包围盒信息\n",
    "        boxes = []\n",
    "        for b in anno['bbox']:\n",
    "            #缩小原始的数值\n",
    "            for i in range(4):\n",
    "                b[i] = b[i]/self.down_scale_factor\n",
    "            \n",
    "            #转换为x0,y0,x1,y1格式\n",
    "            bbox = [b[0],b[1],b[0]+b[2],b[1]+b[3]]\n",
    "            boxes.append(bbox)\n",
    "        \n",
    "        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n",
    "\n",
    "        #获得类别信息，只有1类目标就是指针\n",
    "        num_objs = len(boxes)\n",
    "        labels = torch.ones((num_objs,), dtype=torch.int64)\n",
    "        #获取其他coco格式的必要信息\n",
    "        image_id = torch.tensor([idx])\n",
    "        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n",
    "        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)\n",
    "\n",
    "        #获得标签信息\n",
    "        target = {}\n",
    "        target[\"boxes\"] = boxes\n",
    "        target[\"labels\"] = labels\n",
    "        target[\"image_id\"] = image_id\n",
    "        target[\"area\"] = area\n",
    "        target[\"iscrowd\"] = iscrowd\n",
    "\n",
    "        #进行变化\n",
    "        if self.transforms is not None:\n",
    "            img, target = self.transforms(img, target)\n",
    "        \n",
    "        return img, target        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "def get_transform(train):\n",
    "    transforms = []\n",
    "    transforms.append(T.PILToTensor())\n",
    "    transforms.append(T.ConvertImageDtype(torch.float))\n",
    "    if train:\n",
    "        transforms.append(T.RandomPhotometricDistort())\n",
    "    return T.Compose(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# test the dataset class\n",
    "dataset = MultiPointerMeterDataset(folder, get_transform(True), 6)\n",
    "dataset_test = MultiPointerMeterDataset(folder, get_transform(False), 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "img, anno = dataset.__getitem__(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "img = tT.ToPILImage()(img.to('cpu'))\n",
    "plt.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义目标检测模型\n",
    "使用一个预训练过的faster-rcnn模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=\"DEFAULT\")\n",
    "\n",
    "num_classes = 2  # 1 class (pointer) + background\n",
    "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "from torch import batch_norm_gather_stats_with_counts\n",
    "\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    " dataset, batch_size=2, shuffle=True, num_workers=0,\n",
    " collate_fn=utils.collate_fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# For Training\n",
    "images,targets = next(iter(data_loader))\n",
    "images = list(image for image in images)\n",
    "targets = [{k: v for k, v in t.items()} for t in targets]\n",
    "output = model(images,targets)   # Returns losses and detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# For inference\n",
    "images,targets = next(iter(data_loader))\n",
    "model.eval()\n",
    "# x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]\n",
    "x = list(img.to(torch.device('cuda')) for img in images)\n",
    "predictions = model(x)   # Returns predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 进行训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "n_samples = len(dataset)\n",
    "indices = torch.randperm(n_samples).tolist()\n",
    "dataset = torch.utils.data.Subset(dataset, indices[:-30])\n",
    "dataset_test = torch.utils.data.Subset(dataset_test, indices[-30:])\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=2, shuffle=True, num_workers=0,\n",
    "    collate_fn=utils.collate_fn\n",
    ")\n",
    "\n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_test, batch_size=1,shuffle=False, num_workers=0,\n",
    "    collate_fn=utils.collate_fn\n",
    ")\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "#construct an optimizer\n",
    "params=[p for p in model.parameters() if p.requires_grad]\n",
    "optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9,weight_decay=0.0005)\n",
    "\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# train it\n",
    "num_epochs = 10\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_one_epoch(model, optimizer,data_loader,device, epoch,print_freq=5)\n",
    "    lr_scheduler.step()\n",
    "    evaluate(model,data_loader_test,device=device)\n",
    "\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# save the model to disk\n",
    "# model history:\n",
    "#   9-29 : fasterrcnn_resnet50_fpn，只有指针检测，从头训练10 epoches. 用3P uneven数据，indices[:-30]作为训练集。\n",
    "#   9-30-1 : fasterrcnn_resnet50_fpn，只有指针检测，从头训练10 epoches. 用3P blur数据，indices[:-30]作为训练集。\n",
    "fn = 'model_weights_9-30-1.pth'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), fn) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 查看模型的运行效果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load(fn))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# test the model and show the results\n",
    "\n",
    "# show with pysimplegui\n",
    "%matplotlib auto\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "model.to(device)\n",
    "\n",
    "layout = [  [sg.Text('共有'+str(len(dataset_test))+\"个带标注的水表图像\")],\n",
    "            [sg.Text('请输入要查看的图像序号(从0开始):'), sg.InputText('0')],            \n",
    "            [sg.Button('Ok'), sg.Button('Cancel')] ]\n",
    "\n",
    "# Create the Window\n",
    "window = sg.Window('检测结果查看程序', layout)\n",
    "fig = plt.figure()\n",
    "# Event Loop to process \"events\" and get the \"values\" of the inputs\n",
    "while True:\n",
    "    event, values = window.read()\n",
    "    if event == sg.WIN_CLOSED or event == 'Cancel': # if user closes window or clicks cancel\n",
    "        break\n",
    "    if(int(values[0])>=0 and int(values[0])<len(dataset_test)):\n",
    "        fid = int(values[0])\n",
    "        img,label = dataset_test.__getitem__(fid)\n",
    "        \n",
    "        x = img.unsqueeze(0)\n",
    "        x = x.to(device)\n",
    "        predictions = model(x)  \n",
    "        \n",
    "        plt.cla()\n",
    "        #show the original image\n",
    "        pil_img = tT.ToPILImage()(img)\n",
    "        plt.imshow(pil_img)\n",
    "\n",
    "        #显示包围盒的标签信息\n",
    "        boxes_gt = label['boxes'].numpy()\n",
    "        nbox = boxes_gt.shape[0]\n",
    "        for i in range(nbox):\n",
    "            box = boxes_gt[i,:]\n",
    "            w = box[2]-box[0]\n",
    "            h = box[3]-box[1]\n",
    "            plt.gca().add_patch(plt.Rectangle((box[0],box[1]), w,h,fill=False, color='green'))\n",
    "\n",
    "\n",
    "        #显示包围盒        \n",
    "        boxes = predictions[0]['boxes'].cpu().detach().numpy()\n",
    "        nbox = boxes.shape[0]\n",
    "        for i in range(nbox):\n",
    "            box = boxes[i,:]\n",
    "            w = box[2]-box[0]\n",
    "            h = box[3]-box[1]\n",
    "            plt.gca().add_patch(plt.Rectangle((box[0],box[1]), w,h,fill=False, color='red'))\n",
    "        \n",
    "        \n",
    "        # draw_im_with_anno(folder, flst[fid], anno)\n",
    "        plt.draw()\n",
    "\n",
    "window.close()\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "boxes = predictions[0]['boxes'].cpu().detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "boxes_gt = label['boxes'].numpy()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "myiste",
   "language": "python",
   "display_name": "my_test_jupyter"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "11938c6bc6919ae2720b4d5011047913343b08a43b18698fd82dedb0d4417594"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}